from nl2sql_aline.prompt.value_match_prompt import value_match_prompt
from nl2sql_aline.model.llm import LLM
from nl2sql_aline.tools.extract_block import extract_sql
from nl2sql_aline.model.llm import LLM


class ValueMatch:
    def __init__(self,
                 model: LLM = None) -> None:

        self.prompt = value_match_prompt
        # 给这个类赋予一个属性名为llm，这个属性的数据类型是LLM，或者是None
        self.llm = model

    def sql_reform(self, sql_condition,
                    sql,
                   query,
                   data_info
                   ) -> str:
        """
        用于拆解用户的问题，为子问题
        :param query:
        :return: None or List of String
        """
        # 把Prompt模板，进行填充
        filled_prompt = self.prompt.format(query = query, sql = sql,data_info = data_info, sql_condition = sql_condition)
        # 让大模型对填充好的模板，进行推理
        llm_resp: str = self.llm.chat(prompt=filled_prompt)
        # 对大模型的输出，进行后处理（因为我们没办法保证大模型输出的格式就是我们想要的）
        result = extract_sql(text=llm_resp)
        return result


